#include <iostream>
#include <fstream>
#include <queue>
#include <cassert>
#include <map>
#include <vector>
#include <string>
using namespace std;

#include "Huffman.h"
#include "../BitStream/BitStream.h"

#define END_OF_STREAM 0x400

struct Node
{
    Node *l, *r;
    unsigned short c;
    streamsize freq;

    Node(streamsize frequency = 0, unsigned short symbol = 0, Node *left = NULL, Node *right = NULL)
    {
        freq = frequency;
        c = symbol;
        l = left;
        r = right;
    }

    streamsize calcFreq()
    {
        assert(l && r);
        return (freq = l->freq + r->freq);
    }

    class Comparator
    {
        public:
            bool operator()(const Node *lnode, const Node *rnode) const
            {
                if(lnode->freq > rnode->freq) return true;
                else if(lnode->freq < rnode->freq) return false;
                else return lnode < rnode;  // podle "stari" ukazatele - tohle ne takova blbost...
                // ... mam to tu jen, aby strom odpovidal prikladu, co jsme delali ve skole, ...
                // ... jinak staci jen return lnode->freq > rnode->freq;
            }
    };
};

// projde Huffmanuv strom a zkonstruuje slovnik a zakodovanou reprezentaci stromu
void treeWalk(const Node *node, map<unsigned short,BitStream> &codebook, BitStream &codeword, BitStream &tree_bin, vector<unsigned short> &tree_chars)
{
    if(!node->l || !node->r) // vnitrni uzel Huffmanova stromu ma vzdycky oba potomky
    {	// list stromu
        tree_chars.push_back(node->c);
        codebook[node->c] = codeword;
        return;
    }
    // doleva nuly
    codeword.add(1, 0);
    tree_bin.add(1, 0);
    treeWalk(node->l, codebook, codeword, tree_bin, tree_chars);
    codeword.pop_bit();
    //
    // doprava jednicky
    codeword.add(1, 1);
    tree_bin.add(1, 1);
    treeWalk(node->r, codebook, codeword, tree_bin, tree_chars);
    codeword.pop_bit();
}

BitStream fib2(size_t x)
{
    static vector<size_t> fib;
    static vector<size_t>::iterator it1, it2;
    if(fib.size() == 0)
    {
        fib.push_back(1);
        fib.push_back(1);
        it1 = fib.end() - 2;
        it2 = fib.end() - 1;
    }
    //
    // (do)generovani Fibonacciho posloupnosti, pokud je treba
    while(x >= (*it2))
    {
        fib.push_back((*it1) + (*it2));
        // ++it2; ++it1;    --> bohuzel nejde...zrejme se vlozenim iteratory rozhazi
        it1 = fib.end() - 2;
        it2 = fib.end() - 1;
    }
    vector<size_t>::iterator it = it2 - 1;  // vzdycky jdu nad x...diky tomu nemusim pouzit dalsi podminku
    //
    // kodovani vstupu
    unsigned char nbits = 0;
    unsigned long long val = 0;	// WARNING: 64 bit max!!!
    while(it != fib.begin())
    {
        if(x >= (*it))
        {
            val |= (1ULL << nbits);	// insert it in reverse order
            x -= (*it);
        }	// else (x < (*it)) --> insertion of 0 -- nothing needs to be done (only ++nbits and it's done later for both branches)
        ++nbits;
        --it;
    }
    val = (val << 1) | 1;	// u Fibonacciho kodu se vklada jeste jednicka na konec
    ++nbits;
    //
    BitStream retval;
    while(nbits > 8)	// musim zapisovat po bajtech
        retval.add(8, (unsigned char)(val >> (nbits -= 8)));
    retval.add(nbits, (unsigned char)val);
    //
    return retval;
}

// dekoduje jeden znak puvodniho textu
unsigned short decodeChar(Node *node, unsigned char *code, unsigned int &pos)
{
    if(!node->l) return node->c; // staci otestovat jednoho potomka, protoze Huffmanuv strom ma vzdycky oba nebo zadneho
    //
    if(BitStream::bit_view(code, pos) == 0)
        return decodeChar(node->l, code, ++pos);
    else
        return decodeChar(node->r, code, ++pos);
}

// postavi Huffmanuv strom podle kodu a vrati ukazatel na korenovy uzel
// - jedna se o klasicky rekurzivni sestup
void buildTree(Node *node, const unsigned char *code, unsigned int &pos, queue<unsigned short> &tree_chars)
{
    if(BitStream::bit_view(code, pos) == 0)
    {
        buildTree(node->l = new Node, code, ++pos, tree_chars);
    }
    if(BitStream::bit_view(code, pos) == 1)
    {
        if(!node->l || node->r)
        {
            node->c = tree_chars.front();
            tree_chars.pop();
            return;
        }
        else if(!node->r && (tree_chars.size() == 1))
        {
            node->r = new Node;
            node->r->c = tree_chars.front();
            tree_chars.pop();
            ++pos;
            return;
        }
        buildTree(node->r = new Node, code, ++pos, tree_chars);
    }
}

// dekoduje cele cislo zakodovane Fibbonaciho kodem a vraci ho parameter x
// - navratova hovnota je pozice bezprostredne za koncem zakodovaneho cisla ve vstupnim retezci
unsigned int fib2Dec(const unsigned char *code, unsigned int &x)
{
    static vector<unsigned int> fib;
    static vector<unsigned int>::iterator it1, it2;
    if(fib.size() == 0)
    {
        fib.push_back(1);
        fib.push_back(2);
        it1 = fib.end() - 2;
        it2 = fib.end() - 1;
    }
    // dekodovani
    x = 0;
    size_t pos = 0;
    for(size_t i = 0; ; ++i)
    {
        if(i >= fib.size())
        {   // (do)generovani Fibonacciho posloupnosti, pokud je treba
            fib.push_back((*it1) + (*it2));
            // ++it2; ++it1;    --> bohuzel nejde...zrejme se vlozenim iteratory rozhazi
            it1 = fib.end() - 2;
            it2 = fib.end() - 1;
        }
        //
        if(BitStream::bit_view(code, i) == 1)
        {
            x += fib[i];
            if(BitStream::bit_view(code, i+1) == 1)    // zadne Fibonacciho cislo neobsahuje dve jednicky jinde nez na konci
            {
                pos = i + 1;
                break;
            }
        }
    }
    ++pos;	// indexovano od nuly, takze pocet prectenych bitu je pos+1
    return ((pos / 8) + ((pos % 8) > 0));	// vracim pozici v bytech, takze to musim prepocitat z bitu
}

void
Huffman::Encode(unsigned int /*[in]*/bytes, unsigned short * /*[in]*/data, unsigned char * /*[out]*/encoded, unsigned int &/*[out]*/enc_len)
{
    // Semi-adaptive Huffman encoding
    Node nodes[2*259];	// max 259 leafs: RUNA,RUNB,END_OF_STREAM,256 symbols; 2x: tree has twice as much (minus 1) of all nodes as it has leafs
    streamsize freq[35536];	// frequency table - symbol is 2B value
    int n_nodes;
    priority_queue<Node *, vector<Node *>, Node::Comparator> tree;
    Node *root;
    // encode
    n_nodes = 0;
    memset(freq, 0, 35536*sizeof(streamsize));
    BitStream tree_bin;	// binarni reprezentace Huffmanova stromu
    vector<unsigned short> tree_chars;	// znaky v poradi prochazeni stromu
    BitStream codeword;		// pomocna struktura pro ukladani kodoveho slova
    map<unsigned short,BitStream> codebook;  // tabulka kodovych slov <T,C>
    if(bytes > 0)
    {
        // add END_OF_STREAM symbol
        data[bytes++] = END_OF_STREAM;
        // calc frequency table
        for(streamsize i = 0; i < bytes; i++)
            freq[data[i]]++;
        // fill nodes and insert them into the priority queue
        for(unsigned int i = 0; i < 35536; i++)
        {
            if(freq[i] > 0)
            {
                nodes[n_nodes].c = i;
                nodes[n_nodes].freq = freq[i];
                tree.push(nodes + n_nodes);
                ++n_nodes;
            }
        }
        // build tree
        while(tree.size() > 1) // dokud nezbyde jen koren
        {	// vytvorim rodice a vlozim
            nodes[n_nodes].l = tree.top(); tree.pop();
            nodes[n_nodes].r = tree.top(); tree.pop();
            nodes[n_nodes].calcFreq();
            tree.push(nodes + n_nodes);
            ++n_nodes;
        }
        root = tree.top(); tree.pop();    // ve fronte zbyl koren Huffmanova stromu, ktery si ulozim; fronta je ted prazdna a strom lze traversovat z rootu
        // vyplneni tabulky kodovych slov + vytvoreni binarni reprezentace Huffmanova stromu
        treeWalk(root, codebook, codeword, tree_bin, tree_chars);
        //
        // zakodovani modelu
        const unsigned char *buf;
        enc_len = 0;
        //   1. pocet znaku
        BitStream stream = fib2(codebook.size());
        buf = stream.data();
        memcpy(encoded + enc_len, buf, stream.bytes());
        enc_len += stream.bytes();
        //   2. znaky
        for(size_t i = 0, im = tree_chars.size(); i < im; i++)
        {
            encoded[enc_len++] = (unsigned char)(tree_chars[i] >> 8);
            encoded[enc_len++] = (unsigned char)(tree_chars[i]);
        }
        //   3. strom
        buf = tree_bin.data();
        memcpy(encoded + enc_len, buf, tree_bin.bytes());
        enc_len += tree_bin.bytes();
        // 4. zakodovani vstupniho textu
        stream.clear();
        for(unsigned int i = 0; i < bytes; i++)
            stream += codebook[data[i]];
        buf = stream.data();
        memcpy(encoded + enc_len, buf, stream.bytes());
        enc_len += stream.bytes();
    }
}

void
Huffman::Decode(unsigned int /*[in]*/bytes, unsigned char * /*[in]*/data, unsigned short * /*[out]*/decoded, unsigned int &/*[out]*/dec_len)
{
    Node *root = NULL;          // ukazatel na koren stromu
    map<unsigned short,string> codebook;  // tabulka kodovych slov
    vector<unsigned short> codeword;      // pomocna struktura pro ukladani kodoveho slova
    unsigned int pos;                 // pozice ve vstupnich datech
    unsigned int bit_pos;	// pozice ve streamu v bitech / pouziva se jen v urcitych oblastech, takze to neni absolutni pozice
    unsigned int count;               // pocet symbolu zakodovanych v Huffmanove stromu
    queue<unsigned short> tree_chars;     // znaky zakodovane ve stromu
    unsigned short symbol;
    //
    // inicializace
    codebook.clear();
    codeword.clear();
    if(root) delete root;
    pos = 0;
    //
    // precist pocet znaku
    pos = fib2Dec(data, count);
    //
    // precist znaky zakodovane ve stromu
    for(unsigned int i = 0; i < count; i++)
    {
        symbol = data[pos++];
        symbol = (symbol << 8) | data[pos++];
        tree_chars.push(symbol);
    }
    //
    // precist binarni reprezentaci stromu
    bit_pos = 0;
    buildTree(root = new Node, data+pos, bit_pos, tree_chars);
    pos += (bit_pos / 8) + ((bit_pos % 8) > 0);
    //
    // dekodovat text
    dec_len = 0;
    bit_pos = 0;
    while(1)
    {
        symbol = decodeChar(root, data+pos, bit_pos);
        if(symbol == END_OF_STREAM)
            break;
        decoded[dec_len++] = symbol;
    }
}